package srt.tool;

import java.util.ArrayList;
import java.util.List;
import java.util.Set;

import srt.ast.AssertStmt;
import srt.ast.AssumeStmt;
import srt.ast.BinaryExpr;
import srt.ast.BlockStmt;
import srt.ast.DeclRef;
import srt.ast.EmptyStmt;
import srt.ast.Expr;
import srt.ast.HavocStmt;
import srt.ast.IfStmt;
import srt.ast.IntLiteral;
import srt.ast.Invariant;
import srt.ast.Stmt;
import srt.ast.WhileStmt;
import srt.ast.visitor.impl.DefaultVisitor;

public class LoopAbstractionVisitor extends DefaultVisitor {

	// Implementation of verifier mode
	public LoopAbstractionVisitor() {
		super(true);
	}

	// While while statement encountered do the loop abstraction approximation
	@Override
	public Object visit(WhileStmt whileStmt) {

		List<Stmt> statementsStack = new ArrayList<Stmt>();

		Set<String> havocExpresions = null;
		Expr whileCond = null;

		try {
			whileCond = (BinaryExpr) whileStmt.getCondition();
		} catch (ClassCastException e) {
			whileCond = (DeclRef) whileStmt.getCondition();
		}
		
		//collecting of variable names form while body
		CollectConstraintsVisitor modset = new CollectConstraintsVisitor();
		modset.visit((BlockStmt) whileStmt.getBody());

		havocExpresions = modset.variableNames;
		
		// invariants asserting
		for (Invariant i : whileStmt.getInvariantList().getInvariants()) {
			statementsStack.add(new AssertStmt(i.getExpr()));
		}
		// body variables havocing
		for (String s : havocExpresions) {
			statementsStack.add(new HavocStmt(new DeclRef(s)));
		}
		// invariants assuming
		for (Invariant i : whileStmt.getInvariantList().getInvariants()) {
			statementsStack.add(new AssumeStmt(i.getExpr()));
		}
		
		// while if body
		List<Stmt> endingIf = new ArrayList<Stmt>();

		endingIf.add((Stmt) visit(whileStmt.getBody()));

		for (Invariant i : whileStmt.getInvariantList().getInvariants()) {
			endingIf.add(new AssertStmt(i.getExpr()));
		}
		// assume(false)
		endingIf.add(new AssumeStmt(new IntLiteral(0)));

		statementsStack.add(new IfStmt(whileCond, new BlockStmt(endingIf),
				new EmptyStmt()));

		return new BlockStmt(statementsStack, whileStmt.getNodeInfo());
	}
}